#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RSAAlgorithm特定的Side-Channel Attack实现
"""

import numpy as np
from typing import Dict, Any, List
from ..base_algorithm_attack import AlgorithmSpecificAttack
from .rsa_models import RSALeakageModels


class RSASpecificAttacks(AlgorithmSpecificAttack):
    """RSA特定的Side-Channel Attack"""
    
    def __init__(self):
        super().__init__("RSA")
        
        # RSAParameters
        self.key_size = 2048  # DefaultRSA-2048
        self.block_size = 256  # 2048位 = 256字节
        self.round_count = 1   # RSA没Has轮的概念
        
        # 新架构：分离的模型
        self.intermediate_functions = RSALeakageModels.get_intermediate_functions()
        self.leakage_models = RSALeakageModels.get_leakage_models()
        self.operation_mapping = RSALeakageModels.get_operation_to_intermediate_mapping()
        
    def get_algorithm_parameters(self) -> Dict[str, Any]:
        """获取RSAAlgorithmParameters"""
        return {
            'algorithm': 'RSA',
            'type': 'asymmetric',
            'key_size': self.key_size,
            'block_size': self.block_size,
            'round_count': self.round_count,
            'word_size': self.key_size
        }
    
    def get_intermediate_functions(self) -> Dict[str, Any]:
        """获取Intermediate valueCalculateFunction"""
        return self.intermediate_functions
    
    def get_leakage_models(self) -> Dict[str, Any]:
        """获取泄漏表示模型"""
        return self.leakage_models
    
    def create_leakage_models(self) -> Dict[str, Any]:
        """CreateRSA特定的泄漏模型（用于基Class兼容）"""
        combined = {}
        combined.update(self.intermediate_functions)
        combined.update(self.leakage_models)
        return combined
    
    def get_target_operations(self) -> List[str]:
        """获取RSA可Attack的目标操作"""
        return RSALeakageModels.get_target_operations()
    
    def get_operation_to_intermediate_mapping(self) -> Dict[str, str]:
        """获取目标操作ToIntermediate valueFunction的Mapping"""
        return self.operation_mapping
    
    def _calculate_intermediate_values(self, target_data: np.ndarray, hypothesis: int, 
                                     target_operation: str) -> np.ndarray:
        """
        CalculateRSA特定的Intermediate value
        
        Args:
            target_data: 目标Data（Plaintext字节）
            hypothesis: False设值（Key字节的False设）
            target_operation: 目标操作
            
        Returns:
            Intermediate valueArray（预测的汉明权重）
        """
        if target_operation == 'modular_exponentiation':
            # 改进的模幂运算泄漏模型
            # 与Power仿True模型Match：intermediate = pow(msg_byte, key_byte, 256)
            results = []
            for td in target_data:
                msg_byte = int(td) & 0xFF  # EnsureYes0-255范围
                key_byte = hypothesis & 0xFF  # Key字节False设
                
                # CalculateIntermediate value（与Power仿True一致）
                if msg_byte > 0:
                    intermediate = pow(msg_byte, key_byte, 256)
                else:
                    intermediate = msg_byte
                
                # Calculate汉明权重（与Power模型一致）
                hw_intermediate = bin(intermediate).count('1')
                hw_key = bin(key_byte).count('1')
                
                # 预测的Power泄漏值（与Power仿True的hw_factor公式对应）
                # Update系数以Match增强的泄漏信号
                leakage = hw_intermediate * 0.5 + hw_key * 0.4
                results.append(leakage)
            
            return np.array(results)
        
        elif target_operation == 'square_and_multiply':
            # 平方乘Algorithm泄漏模型
            # False设hypothesisYes私钥的一位
            bit_value = hypothesis & 1
            modulus = 1024  # 简化的模数
            results = []
            for td in target_data:
                base = int(td) % 100
                current_result = base
                new_result, _ = RSALeakageModels.square_and_multiply_model(
                    base, bit_value, current_result, modulus)
                results.append(new_result)
            return np.array(results)
        
        elif target_operation == 'crt_computation':
            # CRTCalculate泄漏模型
            # 简化Process
            p, q = 17, 19  # 小的素数用于演示
            dp = hypothesis % (p - 1)
            dq = hypothesis % (q - 1)
            results = []
            for td in target_data:
                message = int(td) % (p * q)
                result = RSALeakageModels.crt_model(message, dp, dq, p, q)
                results.append(result)
            return np.array(results)
        
        elif target_operation == 'timing_analysis':
            # 时序Attack泄漏模型
            bit_value = hypothesis & 1
            return np.array([RSALeakageModels.timing_attack_model(
                bit_value, RSALeakageModels.hamming_weight(int(td))) for td in target_data])
        
        elif target_operation == 'power_analysis':
            # Power Analysis泄漏模型
            operation_types = ['square', 'multiply', 'add', 'subtract']
            op_type = operation_types[hypothesis % len(operation_types)]
            return np.array([RSALeakageModels.power_analysis_model(
                op_type, RSALeakageModels.hamming_weight(int(td))) for td in target_data])
        
        else:
            # DefaultUse模乘模型
            return np.array([RSALeakageModels.modular_multiplication_model(
                int(td), hypothesis, 65537) for td in target_data])
    
    def _get_selection_function(self, target_data: np.ndarray, hypothesis: int,
                              target_operation: str) -> np.ndarray:
        """
        获取RSA特定的选择Function（用于DPAAttack）
        
        Args:
            target_data: 目标Data
            hypothesis: False设值
            target_operation: 目标操作
            
        Returns:
            选择FunctionResult（0Or1的Array）
        """
        intermediate_values = self._calculate_intermediate_values(
            target_data, hypothesis, target_operation
        )
        
        if target_operation in ['modular_exponentiation', 'square_and_multiply']:
            # UseResult的最低位作For选择Function
            return intermediate_values & 1
        
        elif target_operation == 'crt_computation':
            # UseCRTResult的汉明权重奇偶性
            hamming_weights = np.array([RSALeakageModels.hamming_weight(val) 
                                      for val in intermediate_values])
            return hamming_weights & 1
        
        elif target_operation == 'timing_analysis':
            # Based on时序的选择Function
            # If时序大于阈值则For1，No则For0
            threshold = np.mean(intermediate_values)
            return (intermediate_values > threshold).astype(int)
        
        else:
            # DefaultUse最低位
            return intermediate_values & 1
    
    def attack_square_and_multiply(self, **kwargs):
        """
        Attack平方乘Algorithm
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        return self.run_cpa_attack(
            target_operation='square_and_multiply',
            **kwargs
        )
    
    def attack_crt_implementation(self, **kwargs):
        """
        AttackCRT实现
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        return self.run_dpa_attack(
            target_operation='crt_computation',
            **kwargs
        )
    
    def timing_attack(self, **kwargs):
        """
        时序Attack
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        return self.run_spa_attack(
            target_operation='timing_analysis',
            **kwargs
        )
    
    def power_analysis_attack(self, **kwargs):
        """
        Power AnalysisAttack
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        return self.run_template_attack(
            target_operation='power_analysis',
            **kwargs
        )
    
    def private_key_recovery_attack(self, key_bit_length: int = 16, **kwargs) -> Dict[str, Any]:
        """
        私钥恢复Attack（简化版）
        
        Args:
            key_bit_length: 要恢复的Key位长度
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        results = {}
        recovered_key_bits = {}
        
        self.logger.info(f"Starting RSA private key recovery attack ({key_bit_length} bits)")
        
        # Attack每一位私钥
        for bit_pos in range(key_bit_length):
            self.logger.info(f"Attacking bit {bit_pos}")
            
            try:
                # UseCPAAttack平方乘Algorithm
                result = self.attack_square_and_multiply(
                    target_byte=bit_pos % 16,  # MappingTo字节位置
                    **kwargs
                )
                
                results[f"bit_{bit_pos}"] = result
                
                if result.success and result.recovered_key_bytes:
                    # 将恢复的字节ConvertFor位
                    recovered_bit = result.recovered_key_bytes[0] & 1
                    recovered_key_bits[bit_pos] = recovered_bit
                    self.logger.info(f"Bit {bit_pos}: {recovered_bit} "
                                   f"(confidence: {result.confidence:.4f})")
                else:
                    self.logger.warning(f"Failed to recover bit {bit_pos}")
            
            except Exception as e:
                self.logger.error(f"Error attacking bit {bit_pos}: {e}")
        
        # 构建恢复的私钥位串
        recovered_key_bits_list = []
        for i in range(key_bit_length):
            if i in recovered_key_bits:
                recovered_key_bits_list.append(recovered_key_bits[i])
            else:
                recovered_key_bits_list.append(0)  # 未恢复的位设For0
        
        # 将位串ConvertForInteger
        recovered_key_int = 0
        for i, bit in enumerate(recovered_key_bits_list):
            recovered_key_int |= (bit << i)
        
        success_rate = len(recovered_key_bits) / key_bit_length
        
        summary = {
            'algorithm': 'RSA',
            'attack_type': 'Private Key Recovery',
            'total_bits': key_bit_length,
            'recovered_bits': len(recovered_key_bits),
            'success_rate': success_rate,
            'recovered_key_bits': recovered_key_bits_list,
            'recovered_key_int': recovered_key_int,
            'recovered_key_hex': f'{recovered_key_int:0{(key_bit_length+3)//4}X}',
            'individual_results': results
        }
        
        self.logger.info(f"RSA key recovery completed: {len(recovered_key_bits)}/{key_bit_length} bits recovered "
                        f"({success_rate:.1%} success rate)")
        
        return summary
    
    def validate_recovered_key(self, recovered_key_int: int, 
                             test_message: bytes, test_signature: bytes,
                             public_key: tuple) -> bool:
        """
        Verify恢复的私钥YesNo正确
        
        Args:
            recovered_key_int: 恢复的私钥Integer
            test_message: Test消息
            test_signature: Test签名
            public_key: 公钥 (n, e)
            
        Returns:
            私钥YesNo正确
        """
        try:
            n, e = public_key
            
            # Use恢复的私钥进行签名Verify
            # 这里简化Process，实际RequireComplete的RSAVerify
            message_int = int.from_bytes(test_message, 'big')
            signature_int = int.from_bytes(test_signature, 'big')
            
            # Verify：signature^e ≡ message (mod n)
            verified_message = pow(signature_int, e, n)
            
            return verified_message == message_int
            
        except Exception as e:
            self.logger.error(f"Key validation failed: {e}")
            return False
